import sys, pygame, logging, json
from PyQt5.QtWidgets import *
from PyQt5.QtGui import QPixmap, QImage, QIcon
from PyQt5.QtCore import QTimer, Qt, QPoint

from protolib.gui_utils.pyqt5_messageboxes import show_critical_messagebox

from engine.project._json import JSONBasedProjectLoader as AnimaticLoader

logger = logging.getLogger(__name__)

class TransformToolbar(QWidget):
    def __init__(self):
        super().__init__()

        self.initUI()

    def initUI(self):
        sizePolicy = QSizePolicy()
        sizePolicy.setVerticalStretch(QSizePolicy.MinimumExpanding)

        self.toolbarLabel = QLabel('Object Transforms:')
        self.toolbarLabel.setSizePolicy(sizePolicy)

        self.rotationLabel = QLabel('Rotation:')
        self.rotationLabel.setSizePolicy(sizePolicy)
        self.rotationEdit = QDoubleSpinBox(self)
        self.rotationEdit.setMinimum(-(2 ** 32))
        self.rotationEdit.setMaximum(2 ** 32)

        self.scaleLabel = QLabel('Scale:')
        self.scaleLabel.setSizePolicy(sizePolicy)
        self.xScaleEdit = QDoubleSpinBox(self)
        self.xScaleEdit.setMinimum(-(2 ** 32))
        self.xScaleEdit.setMaximum(2 ** 32)
        self.yScaleEdit = QDoubleSpinBox(self)
        self.yScaleEdit.setMinimum(-(2 ** 32))
        self.yScaleEdit.setMaximum(2 ** 32)

        self.translationLabel = QLabel('Translation:')
        self.translationLabel.setSizePolicy(sizePolicy)
        self.xTranslationEdit = QDoubleSpinBox(self)
        self.xTranslationEdit.setMinimum(-(2 ** 32))
        self.xTranslationEdit.setMaximum(2 ** 32)
        self.yTranslationEdit = QDoubleSpinBox(self)
        self.yTranslationEdit.setMinimum(-(2 ** 32))
        self.yTranslationEdit.setMaximum(2 ** 32)

        self.applyButton = QPushButton('Apply', self)

        self.toolbarLayout = QVBoxLayout()
        self.toolbarLayout.setContentsMargins(0,0,0,0)
        self.toolbarLayout.addWidget(self.toolbarLabel)
        self.toolbarLayout.addWidget(self.rotationLabel)
        self.toolbarLayout.addWidget(self.rotationEdit)
        self.toolbarLayout.addWidget(self.scaleLabel)
        self.toolbarLayout.addWidget(self.xScaleEdit)
        self.toolbarLayout.addWidget(self.yScaleEdit)
        self.toolbarLayout.addWidget(self.translationLabel)
        self.toolbarLayout.addWidget(self.xTranslationEdit)
        self.toolbarLayout.addWidget(self.yTranslationEdit)
        self.toolbarLayout.addWidget(self.applyButton)

        self.setLayout(self.toolbarLayout)

    def zeroEverything(self):
        self.rotationEdit.setValue(0)
        self.xTranslationEdit.setValue(0)
        self.yTranslationEdit.setValue(0)
        self.xScaleEdit.setValue(0)
        self.yScaleEdit.setValue(0)

    def disableEverything(self):
        self.rotationEdit.setDisabled(True)
        self.xTranslationEdit.setDisabled(True)
        self.yTranslationEdit.setDisabled(True)
        self.xScaleEdit.setDisabled(True)
        self.yScaleEdit.setDisabled(True)
        self.applyButton.setDisabled(True)

    def enableEverything(self):
        self.rotationEdit.setDisabled(False)
        self.xTranslationEdit.setDisabled(False)
        self.yTranslationEdit.setDisabled(False)
        self.xScaleEdit.setDisabled(False)
        self.yScaleEdit.setDisabled(False)
        self.applyButton.setDisabled(False)

    def makeEverythingReadOnly(self):
        self.rotationEdit.setReadOnly(True)
        self.xTranslationEdit.setReadOnly(True)
        self.yTranslationEdit.setReadOnly(True)
        self.xScaleEdit.setReadOnly(True)
        self.yScaleEdit.setReadOnly(True)
        self.applyButton.setDisabled(True)

    def makeEverythingNotReadOnly(self):
        self.rotationEdit.setReadOnly(False)
        self.xTranslationEdit.setReadOnly(False)
        self.yTranslationEdit.setReadOnly(False)
        self.xScaleEdit.setReadOnly(False)
        self.yScaleEdit.setReadOnly(False)
        self.applyButton.setDisabled(False)

class MainWindow(QMainWindow):
    def __init__(self, args):
        super().__init__()

        self.loader = AnimaticLoader(args.file_path)
        self.anim_engine = self.loader.get_engine()

        self.setWindowIcon(QIcon('icon.png'))
        self.setWindowTitle("Animatic Maker")

        #region
        menuBar = self.menuBar()
        #region
        self.fileMenu = QMenu("File", self)
        menuBar.addMenu(self.fileMenu)

        #region
        saveAction = QAction("Save", self)
        saveAction.setShortcut("Ctrl+S")
        saveAction.setStatusTip('Save the project to disk')
        saveAction.triggered.connect(self.save_project)
        self.fileMenu.addAction(saveAction)
        #endregion
        #region
        openAction = QAction("Open", self)
        openAction.setShortcut("Ctrl+O")
        openAction.setStatusTip('Open a project on disk')
        openAction.triggered.connect(self.open_project)
        self.fileMenu.addAction(openAction)
        #endregion
        self.fileMenu.addSeparator()
        #region
        quitAction = QAction("Quit", self)
        quitAction.setShortcut("Ctrl+Q")
        quitAction.setStatusTip('Leave the app')
        quitAction.triggered.connect(self.close)
        self.fileMenu.addAction(quitAction)
        #endregion
        #endregion
        #region
        self.playerMenu = QMenu("Player", self)
        menuBar.addMenu(self.playerMenu)
        #region
        playPauseAction = QAction("Play/Pause", self)
        playPauseAction.setShortcut("Space")
        playPauseAction.setStatusTip('Stop or start animation playback')
        playPauseAction.triggered.connect(self.play_pause_animation)
        self.playerMenu.addAction(playPauseAction)
        #endregion
        #endregion
        #endregion

        # Create Pygame surface
        self.pygame_surface = pygame.Surface(self.loader.geometry)

        #region
        # Create the main widget and set it as central
        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        # Create a QLabel for displaying the Pygame surface
        self.pygame_label = QLabel(self)
        self.pygame_label.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding)
        self.pygame_label.setAlignment(Qt.AlignCenter)
        self.pygame_label.mouseMoveEvent = self.pygame_label_clicked

        # Connect the Pygame timer to the update function
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update_animation)
        self.timer.start(1000 // int(self.loader.frame_rate))

        self.add_keyframe_button = QPushButton(self, text='Add keyframe')
        self.add_keyframe_button.clicked.connect(self.add_keyframe)
        self.add_keyframe_button.setDisabled(True)

        self.frame_spinbox = QSpinBox(self)
        self.frame_spinbox.editingFinished.connect(self.frame_spinbox_finished_editing)
        self.frame_spinbox.setMaximum(self.loader.frame_count)
        self.frame_spinbox.setMinimum(1)

        # Create a slider for controlling the current frame
        self.frame_slider = QSlider(self)
        self.frame_slider.valueChanged.connect(self.update_frame_spinbox)
        self.frame_slider.setOrientation(Qt.Horizontal)
        self.frame_slider.setMaximum(self.loader.frame_count)
        self.frame_slider.setMinimum(1)
        self.frame_slider.sliderPressed.connect(self.frame_slider_pressed)
        self.frame_slider.valueChanged.connect(self.frame_slider_value_chanched)

        toolBarSizePolicy = QSizePolicy()
        toolBarSizePolicy.setVerticalStretch(QSizePolicy.MinimumExpanding)

        self.selected_object_label = QLabel('Selected object: None')
        self.selected_object_label.setSizePolicy(toolBarSizePolicy)
        self.object_transforms_at_keyframe = QLabel('Transforms at keyframe: False')
        self.object_transforms_at_keyframe.setSizePolicy(toolBarSizePolicy)

        #region
        playback_layout = QHBoxLayout()

        # Create buttons for play, pause, and stop
        self.play_button = QPushButton("Play", self)
        self.pause_button = QPushButton("Pause", self)
        self.stop_button = QPushButton("Stop", self)

        # Connect buttons to their respective functions
        self.play_button.clicked.connect(self.play_animation)
        self.pause_button.clicked.connect(self.pause_animation)
        self.stop_button.clicked.connect(self.stop_animation)

        playback_layout.addWidget(self.play_button)
        playback_layout.addWidget(self.pause_button)
        playback_layout.addWidget(self.stop_button)
        #endregion

        self.update_display()

        self.transform_toolbar = TransformToolbar()
        self.transform_toolbar.applyButton.clicked.connect(self.update_selected_obj_keyframe)

        layout = QVBoxLayout(central_widget)
        display_toolbar_layout = QHBoxLayout()
        toolbar_layout = QVBoxLayout()

        display_toolbar_layout.addWidget(self.pygame_label)

        toolbar_spacer = QWidget()

        toolbar_layout.addWidget(self.selected_object_label)
        toolbar_layout.addWidget(self.object_transforms_at_keyframe)
        toolbar_layout.addWidget(self.transform_toolbar)
        toolbar_layout.addWidget(toolbar_spacer)
        toolbar_layout.addWidget(self.add_keyframe_button)
        toolbar_layout.addWidget(self.frame_spinbox)

        display_toolbar_layout.addLayout(toolbar_layout)
        layout.addLayout(display_toolbar_layout)
        layout.addWidget(self.frame_slider)
        layout.addLayout(playback_layout)

        self.anim_player_state = 'playing'

        self.stop_animation()
        #endregion

        self.clicked_anim_obj = None

        self.update_selected_anim_obj_display()

    def closeEvent(self, event):
        print('Good bye!')

    def play_animation(self):
        self.timer.start()
        self.play_button.setDisabled(True)
        self.pause_button.setDisabled(False)
        self.stop_button.setDisabled(False)
        if self.frame_slider.value() >= self.loader.frame_count:
            self.frame_slider.setValue(1)
        self.anim_player_state = 'playing'

    def pause_animation(self):
        self.timer.stop()
        self.play_button.setDisabled(False)
        self.pause_button.setDisabled(True)
        self.stop_button.setDisabled(False)
        self.anim_player_state = 'paused'

    def stop_animation(self):
        self.timer.stop()
        self.frame_slider.setValue(0)
        self.play_button.setDisabled(False)
        self.pause_button.setDisabled(True)
        self.stop_button.setDisabled(True)
        self.anim_player_state = 'stopped'

    def play_pause_animation(self):
        if self.anim_player_state in ['paused', 'stopped']:
            self.play_animation()
        else:
            self.pause_animation()

    def update_animation(self):
        self.frame_slider.setValue(self.frame_slider.value() + 1)
        if self.frame_slider.value() >= self.loader.frame_count:
            self.pause_animation()

    def update_display(self):
        self.update_pygame_surface(self.frame_slider.value() - 1)

        image_rect = self.pygame_surface.get_rect()
        viewport_rect = pygame.rect.Rect(0, 0, self.pygame_label.geometry().width(), self.pygame_label.geometry().height())
        final_rect = image_rect.fit(viewport_rect)

        self.pygame_surface_scale_factor = (
            final_rect.width / image_rect.width,
            final_rect.height / image_rect.height
        )

        # Convert Pygame surface to QImage
        img = QImage(pygame.image.tobytes(self.pygame_surface, 'RGB'), self.pygame_surface.get_width(), self.pygame_surface.get_height(), QImage.Format_RGB888)
        img = img.smoothScaled(final_rect.width, final_rect.height)

        # Convert QImage to QPixmap
        pixmap = QPixmap.fromImage(img)

        # Set the QPixmap to the QLabel
        self.pygame_label.setPixmap(pixmap)

        self.update_pygame_surface(self.frame_slider.value() - 1)

    def update_frame_slider(self):
        self.frame_slider.setValue(self.frame_spinbox.value())

    def update_frame_spinbox(self):
        self.frame_spinbox.setValue(self.frame_slider.value())

    def update_pygame_surface(self, number):
        self.pygame_surface.blit(self.anim_engine.render_frame(number), (0, 0))

    def frame_slider_pressed(self):
        if self.anim_player_state == 'playing':
            self.pause_animation()

    def frame_spinbox_finished_editing(self):
        self.update_frame_slider()
        self.update_display()

    def pygame_label_clicked(self, event):
        pos = self.pygame_label.pos()
        x, y = pos.x(), pos.y()
        x *= self.pygame_surface_scale_factor[0]
        y *= self.pygame_surface_scale_factor[1]
        x += event.pos().x()
        y += event.pos().y()

        self.pygame_surface_scale_factor
        data = self.anim_engine.get_anim_obj_rects_and_images(self.frame_slider.value() - 1, True)
        for i in range(len(data[1])):
            if data[1][i].collidepoint((x, y)):
                obj_id = data[2][i]
                self.clicked_anim_obj = obj_id
                self.update_selected_anim_obj_display()
                return
        self.clicked_anim_obj = None
        self.update_selected_anim_obj_display()

    def update_selected_anim_obj_display(self):
        if self.clicked_anim_obj:
            obj = self.anim_engine.anim_objs[self.clicked_anim_obj]
            self.transform_toolbar.rotationEdit.setValue(obj.rotation)
            self.transform_toolbar.xTranslationEdit.setValue(obj.x)
            self.transform_toolbar.yTranslationEdit.setValue(obj.y)
            self.transform_toolbar.xScaleEdit.setValue(obj.width)
            self.transform_toolbar.yScaleEdit.setValue(obj.height)
            self.selected_object_label.setText(f'Selected object: {self.clicked_anim_obj}')
            self.object_transforms_at_keyframe.setText('Transforms at keyframe: False')
            self.transform_toolbar.enableEverything()
            self.transform_toolbar.makeEverythingReadOnly()
            self.add_keyframe_button.setDisabled(False)
            for j in obj.keyframes:
                if j['time'] == self.frame_slider.value() - 1:
                    self.object_transforms_at_keyframe.setText('Transforms at keyframe: True')
                    self.transform_toolbar.makeEverythingNotReadOnly()
                    self.add_keyframe_button.setDisabled(True)
        else:
            self.transform_toolbar.zeroEverything()
            self.transform_toolbar.disableEverything()
            self.selected_object_label.setText('Selected object: None')
            self.object_transforms_at_keyframe.setText('Transforms at keyframe: False')

    def frame_slider_value_chanched(self):
        self.update_display()
        self.update_selected_anim_obj_display()

    def update_selected_obj_keyframe(self):
        if self.clicked_anim_obj:
            if self.frame_slider.value() == 1:
                self.loader.objects[self.clicked_anim_obj]['position'] = [self.transform_toolbar.xTranslationEdit.value(), self.transform_toolbar.yTranslationEdit.value()]
                self.loader.objects[self.clicked_anim_obj]['geometry'] = [self.transform_toolbar.xScaleEdit.value(), self.transform_toolbar.yScaleEdit.value()]
                self.loader.objects[self.clicked_anim_obj]['rotation'] = self.transform_toolbar.rotationEdit.value()
                self.anim_engine = self.loader.get_engine()
                self.update_selected_anim_obj_display()
                self.update_display()
            else:
                for keyframe, _ in enumerate(self.loader.keyframes):
                    if self.loader.keyframes[keyframe]['time'] == self.frame_slider.value() - 1:
                        self.loader.keyframes[keyframe]['position'] = [self.transform_toolbar.xTranslationEdit.value(), self.transform_toolbar.yTranslationEdit.value()]
                        self.loader.keyframes[keyframe]['geometry'] = [self.transform_toolbar.xScaleEdit.value(), self.transform_toolbar.yScaleEdit.value()]
                        self.loader.keyframes[keyframe]['rotation'] = self.transform_toolbar.rotationEdit.value()
                        self.anim_engine = self.loader.get_engine()
                        self.update_selected_anim_obj_display()
                        self.update_display()
                        break

    def save_project(self):
        options = QFileDialog.Options()
        options |= QFileDialog.DontUseNativeDialog
        dir_path = QFileDialog.getExistingDirectory(self, 'Select directory to save project in')
        if dir_path:
            save_data = self.loader.get_export_data()
            for file_name in save_data:
                file_data = save_data[file_name]
                with open(f'{dir_path}/{file_name}', 'w') as f:
                    json.dump(file_data, f)

    def open_project(self):
        options = QFileDialog.Options()
        options |= QFileDialog.DontUseNativeDialog
        dir_path = QFileDialog.getExistingDirectory(self, 'Select directory containing a project')
        if dir_path:
            self.loader.load_project(dir_path)
            self.anim_engine = self.loader.get_engine()
            self.frame_slider.setValue(0)

    def add_keyframe(self):
        if self.clicked_anim_obj:
            obj = self.anim_engine.anim_objs[self.clicked_anim_obj]
            on_keyframe = False
            for i, j in enumerate(obj.keyframes):
                if j['time'] == self.frame_slider.value() - 1:
                    on_keyframe = True
                    break
            if not on_keyframe:
                self.loader.add_keyframe(
                    self.frame_slider.value() - 1,
                    self.transform_toolbar.xTranslationEdit.value(),
                    self.transform_toolbar.yTranslationEdit.value(),
                    self.transform_toolbar.xScaleEdit.value(),
                    self.transform_toolbar.yScaleEdit.value(),
                    self.transform_toolbar.rotationEdit.value(),
                    self.clicked_anim_obj,
                    'main.json' # Temporary
                )
                self.anim_engine = self.loader.get_engine()
            self.update_selected_anim_obj_display()

    def resizeEvent(self, event):
        QMainWindow.resizeEvent(self, event)
        self.update_display()

def main(args):
    app = QApplication(sys.argv)
    if not args.file_path:
        logger.critical('A project path must be specified')
        show_critical_messagebox(message="A project path must be specified")
        sys.exit(1)
    window = MainWindow(args)
    window.show()
    app.exec()
